import json
import asyncio
import aiofiles
from functools import partial

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import time

from starlette.responses import StreamingResponse

from LLM_Prompt.prompt_workflow import classify_and_extract
from oracle_api.conn import connect_to_oracle
from boat_nl2sql.utiles import extract_sql_from_markdown
from oracle_api.dispatch_Dataapi import call_api
from ollama_stream.sql_analyse_output import analyze_boat_flow_api_data
from ollama_stream.sql_output import sql_general
from LLM_Prompt.prompt_echarts import visualize_sql_result


async def async_analyze_boat_flow_api_data(question, sql_result):
    for item in analyze_boat_flow_api_data(question, sql_result):
        yield item
        await asyncio.sleep(0)  # 避免阻塞事件循环


async def generate_html(sql_result, extracted_sql, question):
    try:
        text2sql_result = []
        for item in sql_result:
            name = None
            count = None
            for k, v in item.items():
                if isinstance(v, str):
                    if v.isdigit():
                        count = int(v)
                    else:
                        name = v
            text2sql_result.append((name, count))

        loop = asyncio.get_event_loop()
        output_file = await loop.run_in_executor(
            None,
            partial(visualize_sql_result, text2sql_result, extracted_sql, question)
        )

        if output_file:
            print(f"可视化已完成，请打开以下文件查看: {output_file}")
            try:
                async with aiofiles.open(output_file, 'r', encoding='utf-8') as f:
                    return await f.read()
            except Exception as e:
                return f"无法读取HTML文件: {str(e)}"
    except Exception as e:
        print(f"执行查询或可视化时出错: {e}")
        return None


async def process_sql_and_visualize(question, sql_result, extracted_sql, action_name):
    html_task = None
    if len(sql_result) > 1 and action_name == "not matched":
        html_task = asyncio.create_task(generate_html(sql_result, extracted_sql, question))

    async for content in async_analyze_boat_flow_api_data(question, sql_result):
        yield f"data: {json.dumps({'type': 'content', 'content': content})}\n\n"
        print(content, end="", flush=True)

    if html_task:
        html_content = await html_task
        if html_content:
            yield f"data: {json.dumps({'type': 'html', 'content': html_content})}\n\n"

app = FastAPI()

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

API_STRATEGIES = [
    "vessel_trajectory", "checkpoint_passage", "violation_detection", "checkpoint_traffic"
]


class QuestionRequest(BaseModel):
    question: str


@app.post("/LLM")
async def query_handler(request: QuestionRequest):
    question = request.question
    if not question:
        raise HTTPException(status_code=400, detail="缺少问题参数")

    # 问题分类+参数提取
    start = time.time()
    classify_prompt = classify_and_extract(question)
    classify_time = time.time() - start
    print(f"问题分类+参数提取花费时间：{classify_time:.4f}秒")
    print(classify_prompt)

    request_dict = json.loads(classify_prompt)
    data = call_api(request_dict)
    action_name, payload = data.get('api', None), data.get('payload', None)

    # 流式响应生成器
    async def generate():
        global extracted_sql
        try:
            if action_name in API_STRATEGIES:
                sql_result = data.get('data', None)
                # 构造响应数据
                response_data = {
                    'type': 'result',
                    "result": {
                        "action_name": action_name,
                        "params": payload
                    }
                }
                # 生成SSE格式的响应
                yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"

                # sql结果分析
                for content in analyze_boat_flow_api_data(question, sql_result):
                    yield f"data: {json.dumps({'type': 'content', 'content': content})}\n\n"
                    print(content, end="", flush=True)  # 外部处理实时输出

            else:
                accumulated_content = []
                for content in sql_general(question):
                    print(content, end="", flush=True)  # 外部处理实时输出
                    accumulated_content.append(content)
                    yield f"data: {json.dumps({'type': 'text2sql', 'content': content})}\n\n"

                # 合并所有内容为完整SQL
                full_sql = ''.join(accumulated_content)

                # 提取SQL（假设存在Markdown代码块）
                if "```sql" in full_sql:
                    extracted_sql = extract_sql_from_markdown(full_sql)
                else:
                    extracted_sql = full_sql.strip().strip(';')  # 去除多余分号

                print("\nsql语句：")
                print(extracted_sql)

                # 执行SQL并获取结果
                try:
                    # 获取数据库连接和游标
                    conn = connect_to_oracle()
                    cursor = conn.cursor()
                    # 执行 SQL 查询，传入参数作为位置绑定
                    cursor.execute(extracted_sql)
                except Exception as e:
                    error_msg = f"SQL执行失败: {str(e)}"
                    yield f"data: {json.dumps({'type': 'error', 'error': error_msg})}\n\n"
                    return  # 终止生成器

                # 获取查询结果的列名
                columns = [col[0] for col in cursor.description]
                sql_result = []
                for row in cursor.fetchall():
                    # 将查询结果转为字典，并转换键值为 Python 原生类型
                    row_dict = dict(zip(columns, row))
                    # 强制将所有键和值转换为 Python 的原生类型（字符串、整数、浮点数等）
                    row_dict = {str(key): (str(value) if value is not None else None) for key, value in
                                row_dict.items()}
                    sql_result.append(row_dict)

                print("查询结果：")
                print(sql_result)

                response_data = {
                    'type': 'result',
                    "result": {
                        "action_name": action_name,
                        "params": payload
                    }
                }

                print("返回result：")
                print(response_data)

                # 生成SSE格式的响应
                yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"

                # sql结果分析 与 html图表生成异步执行
                async for item in process_sql_and_visualize(question, sql_result, extracted_sql, action_name):
                    yield item

        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'error': str(e)})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=3005)
